/*
 *  Copyright 2010 INFN - APE group
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <cuos/cuos.h>

//----------------------------------------------------------------------

// Communicators
typedef int cuMPI_Comm;
#define cuMPI_COMM_NULL                 0
#define cuMPI_COMM_WORLD                1
#define cuMPI_COMM_SELF                 2
#define cuMPI_NUM_COMMS                 3

// Data types
typedef int cuMPI_Datatype;
#define cuMPI_DATATYPE_NULL             0
#define cuMPI_BYTE                      1
#define cuMPI_PACKED                    2
#define cuMPI_CHAR                      3
#define cuMPI_SHORT                     4
#define cuMPI_INT                       5
#define cuMPI_LONG                      6
#define cuMPI_FLOAT                     7
#define cuMPI_DOUBLE                    8
#define cuMPI_NUM_TYPES                 9


// Collective operations
typedef int cuMPI_Op;
#define cuMPI_MAX                       0
#define cuMPI_MIN                       1
#define cuMPI_SUM                       2
#define cuMPI_PROD                      3
#define cuMPI_NUM_OPS                   4

// Error codes
#define cuMPI_SUCCESS                   0  
#define cuMPI_ERR_BUFFER                1
#define cuMPI_ERR_COUNT                 2
#define cuMPI_ERR_TYPE                  3
#define cuMPI_ERR_TAG                   4
#define cuMPI_ERR_COMM                  5
#define cuMPI_ERR_RANK                  6
#define cuMPI_ERR_REQUEST               7
#define cuMPI_ERR_ROOT                  8
#define cuMPI_ERR_GROUP                 9
#define cuMPI_ERR_OP                    10
#define cuMPI_ERR_TOPOLOGY              11
#define cuMPI_ERR_DIMS                  12
#define cuMPI_ERR_ARG                   13
#define cuMPI_ERR_UNKNOWN               14

// Non-blocking requests
struct cumpi_request {
    int mpi_error; // MPI_SUCCESS or any other error
    enum {
        n_flags = 30 // good on C1060
    };
    int flags[n_flags]; // >0 if completed, 0 otherwise
};
typedef struct cumpi_request cuMPI_Request;

#define INIT_CUMPIREQUEST(B) { cuMPI_SUCCESS, { B,B,B,B,B, B,B,B,B,B,  B,B,B,B,B, B,B,B,B,B,  B,B,B,B,B, B,B,B,B,B } }

//----------------------------------------------------------------------

__device__ cuos::sys_req_event_t gbl_req_event; // = INIT_GPU_EVENT(CUOS_INIT_DESC);
__device__ cuos::sys_rep_event_t gbl_rep_event;

//using namespace cuos::device;

//----------------------------------------------------------------------

static __device__ int cuMPI_Send(cuos::device::sys_queue_t *sys_q, void *glb_buf, int count, cuMPI_Datatype datatype, int dest, int tag, cuMPI_Comm comm)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    const uint block_size = blockDim.x*blockDim.y*blockDim.z;
    //const uint grid_size = gridDim.y*gridDim.x;
    const uint gid = tid + bid*block_size;
    int retcode = cuMPI_SUCCESS;

    using namespace cuos::system_service;

    cuos::sys_req_t sys_desc;
    sys_desc.id = sys_mpi_send;
    sys_desc.serial = 0;
    sys_desc.n_params = 6;
    sys_desc.flags = request::flag_pending|request::flag_need_reply;
    sys_desc.params[0] = cuos_make_ptr_param( glb_buf, 1 );
    sys_desc.params[1] = cuos_make_int_param( count );
    sys_desc.params[2] = cuos_make_int_param( datatype );
    sys_desc.params[3] = cuos_make_int_param( dest );
    sys_desc.params[4] = cuos_make_int_param( tag );
    sys_desc.params[5] = cuos_make_int_param( comm );
    sys_desc.params[6] = cuos_make_none_param();
    sys_desc.params[7] = cuos_make_none_param();
    if(gid) {
        gbl_req_event.u.t = sys_desc;
    }
    __syncthreads();

    int ret = cuos::device::syscall_exec(sys_q, &gbl_req_event, &gbl_rep_event);
    if(ret != cuos::SUCCESS) {
        cuos_error("[%d:%d] error %d in syscall_exec()\n", bid, tid, ret);
        retcode = cuMPI_ERR_UNKNOWN;
    } else
        retcode = gbl_rep_event.u.t.retval.si;
    
    return retcode;
}

//----------------------------------------------------------------------

static __device__ int cuMPI_Recv(cuos::device::sys_queue_t *sys_q, void *glb_buf, int count, cuMPI_Datatype datatype, int source, int tag, cuMPI_Comm comm)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    const uint block_size = blockDim.x*blockDim.y*blockDim.z;
    //const uint grid_size = gridDim.y*gridDim.x;
    const uint gid = tid + bid*block_size;
    int retcode = cuMPI_SUCCESS;

    using namespace cuos::system_service;

    cuos::sys_req_t sys_desc;
        
    sys_desc.id = sys_mpi_recv;
    sys_desc.serial = 0;
    sys_desc.n_params = 6;
    sys_desc.flags = request::flag_pending|request::flag_need_reply;
    sys_desc.params[0] = cuos_make_ptr_param( glb_buf, 1 );
    sys_desc.params[1] = cuos_make_int_param( count );
    sys_desc.params[2] = cuos_make_int_param( datatype );
    sys_desc.params[3] = cuos_make_int_param( source );
    sys_desc.params[4] = cuos_make_int_param( tag );
    sys_desc.params[5] = cuos_make_int_param( comm );
    sys_desc.params[6] = cuos_make_none_param();
    sys_desc.params[7] = cuos_make_none_param();
    if(bid == 0 && tid == 0) {
        gbl_req_event.u.t = sys_desc;
    }
    __syncthreads();

    int ret = cuos::device::syscall_exec(sys_q, &gbl_req_event, &gbl_rep_event);
    if(ret != cuos::SUCCESS) {
        cuos_error("[%d:%d] error %d in syscall_exec()\n", bid, tid, ret);
        retcode = cuMPI_ERR_UNKNOWN;
    } else
        retcode = gbl_rep_event.u.t.retval.si;
    
    return retcode;
}

//----------------------------------------------------------------------

static __device__ int cuMPI_Reduce(cuos::device::sys_queue_t *sys_q, void *glb_sendbuf, void *glb_recvbuf, int count, cuMPI_Datatype datatype, cuMPI_Op op, int root, cuMPI_Comm comm)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    const uint block_size = blockDim.x*blockDim.y*blockDim.z;
    //const uint grid_size = gridDim.y*gridDim.x;
    const uint gid = tid + bid*block_size;
    int retcode = cuMPI_SUCCESS;

    using namespace cuos::system_service;

    if(gid == 0) {
        cuos::sys_req_t sys_desc;

        sys_desc.id = sys_mpi_reduce;
        sys_desc.serial = 0;
        sys_desc.n_params = 7;
        sys_desc.flags = request::flag_pending|request::flag_need_reply;
        sys_desc.params[0] = cuos_make_ptr_param( glb_sendbuf, 1 );
        sys_desc.params[1] = cuos_make_ptr_param( glb_recvbuf, 1 );
        sys_desc.params[2] = cuos_make_int_param( count );
        sys_desc.params[3] = cuos_make_int_param( datatype );
        sys_desc.params[4] = cuos_make_int_param( op );
        sys_desc.params[5] = cuos_make_int_param( root );
        sys_desc.params[6] = cuos_make_int_param( comm );
        sys_desc.params[7] = cuos_make_none_param();
        //sys_desc.retval = cuos_make_int_param( 0 );

        gbl_req_event.u.t = sys_desc;
    }
    __syncthreads();

    int ret = cuos::device::syscall_exec(sys_q, &gbl_req_event, &gbl_rep_event);
    if(ret != cuos::SUCCESS) {
        cuos_error("error in syscall_exec()\n");
        retcode = cuMPI_ERR_UNKNOWN;
    }
    // don't wait for return
    return retcode;
}

//----------------------------------------------------------------------

static __device__ int cuMPI_Isend(cuos::device::sys_queue_t *sys_q, void *glb_buf, int count, cuMPI_Datatype datatype, int dest, int tag, cuMPI_Comm comm, cuMPI_Request *glb_request)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    const uint block_size = blockDim.x*blockDim.y*blockDim.z;
    //const uint grid_size = gridDim.y*gridDim.x;
    const uint gid = tid + bid*block_size;
    int retcode = cuMPI_SUCCESS;

    using namespace cuos::system_service;

    if(gid == 0) {
        cuos::sys_req_t sys_desc;

        sys_desc.id = sys_mpi_isend;
        sys_desc.serial = 0;
        sys_desc.n_params = 7;
        sys_desc.flags = request::flag_pending|request::flag_request;
        sys_desc.params[0] = cuos_make_ptr_param( glb_buf, 1 );
        sys_desc.params[1] = cuos_make_int_param( count );
        sys_desc.params[2] = cuos_make_int_param( datatype );
        sys_desc.params[3] = cuos_make_int_param( dest );
        sys_desc.params[4] = cuos_make_int_param( tag );
        sys_desc.params[5] = cuos_make_int_param( comm );
        sys_desc.params[6] = cuos_make_ptr_param( glb_request, 1 );
        sys_desc.params[7] = cuos_make_none_param();
        //sys_desc.retval = cuos_make_int_param( 0 );

        gbl_req_event.u.t = sys_desc;

        *(volatile int*)&glb_request->mpi_error = cuMPI_SUCCESS;
    }
    if(tid == 0) {
        *(volatile int*)&glb_request->flags[bid] = 0;
        __threadfence();
    }
    __syncthreads();

    int ret = cuos::device::syscall_exec_oneway(sys_q, &gbl_req_event);
    if(ret != cuos::SUCCESS) {
        cuos_error("error in syscall_exec()\n");
        retcode = cuMPI_ERR_UNKNOWN;
    }
    // don't wait for return
    return retcode;
}

//----------------------------------------------------------------------

static __device__ int cuMPI_Irecv(cuos::device::sys_queue_t *sys_q, void *glb_buf, int count, cuMPI_Datatype datatype, int source, int tag, cuMPI_Comm comm, cuMPI_Request *glb_request)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    const uint block_size = blockDim.x*blockDim.y*blockDim.z;
    //const uint grid_size = gridDim.y*gridDim.x;
    const uint gid = tid + bid*block_size;
    int retcode = cuMPI_SUCCESS;

    using namespace cuos::system_service;

    if(gid == 0) {
        cuos::sys_req_t sys_desc;

        sys_desc.id = sys_mpi_irecv;
        sys_desc.serial = 0;
        sys_desc.n_params = 7;
        sys_desc.flags = request::flag_pending|request::flag_request;
        sys_desc.params[0] = cuos_make_ptr_param( glb_buf, 1 );
        sys_desc.params[1] = cuos_make_int_param( count );
        sys_desc.params[2] = cuos_make_int_param( datatype );
        sys_desc.params[3] = cuos_make_int_param( source );
        sys_desc.params[4] = cuos_make_int_param( tag );
        sys_desc.params[5] = cuos_make_int_param( comm );
        sys_desc.params[6] = cuos_make_ptr_param( glb_request, 1 );
        sys_desc.params[7] = cuos_make_none_param();
        //sys_desc.retval = cuos_make_int_param( 0 );

        gbl_req_event.u.t = sys_desc;

        *(volatile int*)&glb_request->mpi_error = cuMPI_SUCCESS;
    }
    if(tid == 0) {
        *(volatile int*)&glb_request->flags[bid] = 0;
        __threadfence();
    }
    __syncthreads();

    int ret = cuos::device::syscall_exec_oneway(sys_q, &gbl_req_event);
    if(ret != cuos::SUCCESS) {
        cuos_error("error in syscall_exec()\n");
        retcode = cuMPI_ERR_UNKNOWN;
    }
    // don't wait for return
    return retcode;
}

//----------------------------------------------------------------------

/*! Tests for the completion of a specific \c request (send or receive).
 *
 *  \p glb_request The \c request of interest.
 *  \return \c glb_request state: > 0 if request is done,
 *  cuMPI_SUCCESS if not done or any other MPI error code.
 */

static inline __device__ int cuos_atomic_read(int &i)
{
    return *(volatile int*)&(i);
}

static __device__ int cuMPI_Test(cuos::device::sys_queue_t *sys_q, cuMPI_Request *glb_request)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    int ret = cuMPI_SUCCESS;
    if(*(volatile int*)&glb_request->mpi_error != cuMPI_SUCCESS)
        ret = glb_request->mpi_error;
    else {
        int flag = *(volatile int*)&glb_request->flags[bid];
        ret = flag; //!!flag;
        if(flag > 0) {
            __syncthreads(); // sync all threads before resetting flag
            if(tid == 0) {
                *(volatile int*)&glb_request->flags[bid] = 0;
                __threadfence();
            }
        }
    }
    __syncthreads();

    return ret;
}

//----------------------------------------------------------------------

/*! Waits for the completion of a specific \c request (send or receive).
 *
 *  \p glb_request The \c request of interest.
 *  \return \c glb_request state: > 0 if request is done,
 *  cuMPI_SUCCESS if not done or any other MPI error code.
 */

static __device__ int cuMPI_Wait(cuos::device::sys_queue_t *sys_q, cuMPI_Request *glb_request)
{
    const uint tid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    const uint bid = blockIdx.x + gridDim.x*blockIdx.y;
    int ret = cuMPI_SUCCESS;
    while(true) {
        if(*(volatile int*)&glb_request->mpi_error != cuMPI_SUCCESS) {
            ret = glb_request->mpi_error;
            break;
        }
        int flag = *(volatile int*)&glb_request->flags[bid];
        ret = flag; //!!flag;
        if(flag > 0) {
            __syncthreads(); // sync all threads before resetting flag
            if(tid == 0) {
                *(volatile int*)&glb_request->flags[bid] = 0;
                __threadfence();
            }
            break;
        }
#if 0
        int mpi_error = *(volatile int*)&glb_request->mpi_error;
        if(mpi_error != cuMPI_SUCCESS) {
            ret = glb_request->mpi_error;
            break;
        }
        bool completed = !!*(volatile int*)&glb_request->req_completed;
        if(completed) {
            ret = 1;
            break;
        }
#endif
    }
    __syncthreads();

    return ret;
}

//----------------------------------------------------------------------

/*! Waits for the completion of a specific \c request (send or receive).
 *
 *  \p glb_request The \c request of interest.
 *  \return \c glb_request state: > 0 if request is done,
 *  cuMPI_SUCCESS if not done or any other MPI error code.
 */

static __device__ int cuMPI_Waitall(cuos::device::sys_queue_t *sys_q, int count, cuMPI_Request *glb_array_of_requests)
{
    int ret = cuMPI_SUCCESS;
    for(int i = 0; i<count; ++i) {
        ret = cuMPI_Wait(sys_q, glb_array_of_requests + i);
        if(ret < cuMPI_SUCCESS)
            break;
    }
    return ret;
}

//----------------------------------------------------------------------

/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */
